//------------------------------------------------------------------------------
// RegPropertyStore.cpp
//------------------------------------------------------------------------------
// File provided for Microsoft Shared Source.
// Copyright (c) Microsoft Corporation. All rights reserved.
//
// THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF 
// ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO 
// THE IMPLIED WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A 
// PARTICULAR PURPOSE.
//------------------------------------------------------------------------------
//
// Property store implementation for reg folder
//
//------------------------------------------------------------------------------

#include "pch.h"
#include "RegNamespace.h"

//------------------------------------------------------------------------------
// CRegPropertyStore
//------------------------------------------------------------------------------

class CRegPropertyStore :
    public CComObjectRootEx<CComSingleThreadModel>,
    public IPropertyStore
{
public:
    CRegPropertyStore() : _flags(GPS_DEFAULT)
    {
    }

    BEGIN_COM_MAP(CRegPropertyStore)
        COM_INTERFACE_ENTRY(IPropertyStore)
    END_COM_MAP()

    DECLARE_PROTECT_FINAL_CONSTRUCT()

    static HRESULT s_CreateInstance(const PROPERTYKEY *rgKeys, UINT cKeys, GETPROPERTYSTOREFLAGS flags, __in IShellFolder *psf, PCUITEMID_CHILD pidl, REFIID riid, __deref_out void **ppv);

    HRESULT FinalConstruct()
    {
        return S_OK;
    }

    void FinalRelease()
    {
    }

    // IPropertyStore
    IFACEMETHODIMP GetCount(__out DWORD *cProps);
    IFACEMETHODIMP GetAt(DWORD iProp, __out PROPERTYKEY *pkey);
    IFACEMETHODIMP GetValue(REFPROPERTYKEY key, __out PROPVARIANT *ppropvar);
    IFACEMETHODIMP SetValue(REFPROPERTYKEY key, REFPROPVARIANT propvar);
    IFACEMETHODIMP Commit();

private:
    HRESULT _Initialize(const PROPERTYKEY *rgKeys, UINT cKeys, GETPROPERTYSTOREFLAGS flags, __in IShellFolder *psf, PCUITEMID_CHILD pidl);
    HRESULT _InitializePropertyStore(const PROPERTYKEY *rgKeys, UINT cKeys, GETPROPERTYSTOREFLAGS flags);

private:
    CComPtr<IShellFolder2> _spsf;
    CComMemPtr<ITEMID_CHILD> _spidl;
    CComPtr<IPropertyStore> _sppropstore;
    GETPROPERTYSTOREFLAGS _flags;
    CAtlArray<PROPERTYKEY> _rgDelayedKeys;
    CAtlArray<PROPERTYKEY> _rgDirty;
};

HRESULT CRegPropertyStore::_InitializePropertyStore(const PROPERTYKEY *rgKeys, UINT cKeys, GETPROPERTYSTOREFLAGS flags)
{
    HRESULT hr = PSCreateMemoryPropertyStore(IID_PPV_ARGS(&_sppropstore));
    if (SUCCEEDED(hr))
    {
        _flags = flags;
        CComPtr<IRegPropertyHandler> sprph;
        hr = _spsf->QueryInterface(IID_PPV_ARGS(&sprph));
        for (UINT i = 0; i < cKeys && SUCCEEDED(hr); i++)
        {
            if (flags & GPS_DELAYCREATION)
            {
                _rgDelayedKeys.Add(rgKeys[i]);
            }
            else
            {
                CComPropVariant spropvar;
                hr = sprph->GetProperty(_spidl, rgKeys[i], &spropvar);
                if (FAILED(hr))
                {
                    // Intentional stomping of the previous HRESULT 
                    // 
                    // VT_EMPTY means that the value is not available for this item, so clear the 
                    // propvariant and insert an empty instead
                    hr = spropvar.Clear();
                }

                if (SUCCEEDED(hr))
                {
                    // It's ok for the property to not be there, the PROPVARIANT will be VT_EMPTY
                    hr = _sppropstore->SetValue(rgKeys[i], spropvar);
                }
            }
        }
    }
    return hr;
}

HRESULT CRegPropertyStore::_Initialize(const PROPERTYKEY *rgKeys, UINT cKeys, GETPROPERTYSTOREFLAGS flags, __in IShellFolder *psf, PCUITEMID_CHILD pidl)
{
    HRESULT hr = psf->QueryInterface(IID_PPV_ARGS(&_spsf));
    if (SUCCEEDED(hr))
    {
        hr = CloneChildIDList(pidl, &_spidl);
        if (SUCCEEDED(hr))
        {
            hr = _InitializePropertyStore(rgKeys, cKeys, flags);
        }
    }
    return hr;
}

HRESULT CRegPropertyStore::s_CreateInstance(const PROPERTYKEY *rgKeys, UINT cKeys, GETPROPERTYSTOREFLAGS flags, __in IShellFolder *psf, PCUITEMID_CHILD pidl, REFIID riid, __deref_out void **ppv)
{
    *ppv = NULL;
    CComObjectWithRef<CRegPropertyStore> *psprps;
    HRESULT hr = CComObjectWithRef<CRegPropertyStore>::CreateInstance(&psprps);
    if (SUCCEEDED(hr))
    {
        hr = psprps->_Initialize(rgKeys, cKeys, flags, psf, pidl);
        if (SUCCEEDED(hr))
        {
            hr = psprps->QueryInterface(riid, ppv);
        }
        psprps->Release();
    }
    return hr;
}

// IPropertyStore
STDMETHODIMP CRegPropertyStore::GetCount(__out DWORD *cProps)
{
    return _sppropstore->GetCount(cProps);
}

STDMETHODIMP CRegPropertyStore::GetAt(DWORD iProp, __out PROPERTYKEY *pkey)
{
    return _sppropstore->GetAt(iProp, pkey);
}

STDMETHODIMP CRegPropertyStore::GetValue(REFPROPERTYKEY key, __out PROPVARIANT *ppropvar)
{
    HRESULT hr = S_OK;
    if (_flags & GPS_DELAYCREATION)
    {
        // Find this prop in the delayed key list
        for (size_t i = 0; i < _rgDelayedKeys.GetCount() && SUCCEEDED(hr); i++)
        {
            if (IsEqualPropertyKey(key, _rgDelayedKeys[i]))
            {
                _rgDelayedKeys.RemoveAt(i);
                CComPtr<IRegPropertyHandler> sprph;
                hr = _spsf->QueryInterface(IID_PPV_ARGS(&sprph));
                if (SUCCEEDED(hr))
                {
                    CComPropVariant spropvar;
                    hr = sprph->GetProperty(_spidl, key, &spropvar);
                    if (SUCCEEDED(hr) || hr == E_INVALIDARG)
                    {
                        hr = _sppropstore->SetValue(key, spropvar);
                        break;
                    }
                }
            }
        }
    }
    if (SUCCEEDED(hr))
    {
        hr = _sppropstore->GetValue(key, ppropvar);
    }
    return hr;
}

STDMETHODIMP CRegPropertyStore::SetValue(REFPROPERTYKEY key, REFPROPVARIANT propvar)
{
    HRESULT hr = _sppropstore->SetValue(key, propvar);
    if (SUCCEEDED(hr))
    {
        try
        {
            _rgDirty.Add(key);
        }
        catch (CAtlException &e)
        {
            hr = e.m_hr;
        }
    }
    return hr;
}

STDMETHODIMP CRegPropertyStore::Commit()
{
    CComPtr<IRegPropertyHandler> sprph;
    HRESULT hr = _spsf->QueryInterface(IID_PPV_ARGS(&sprph));
    for (size_t i = 0; i < _rgDirty.GetCount() && SUCCEEDED(hr); i++)
    {
        CComPropVariant sppropvar;
        hr = _sppropstore->GetValue(_rgDirty.GetAt(i), &sppropvar);
        if (SUCCEEDED(hr))
        {
            hr = sprph->SetProperty(_spidl, _rgDirty.GetAt(i), sppropvar);
        }
    }
    return hr;
}

//------------------------------------------------------------------------------
// CRegPropertyStoreFactory
//------------------------------------------------------------------------------

class CRegPropertyStoreFactory : 
    public CComObjectRootEx<CComSingleThreadModel>,
    public IPropertyStoreFactory
{
public:
    CRegPropertyStoreFactory()
    {
    }

    BEGIN_COM_MAP(CRegPropertyStoreFactory)
        COM_INTERFACE_ENTRY(IPropertyStoreFactory)
    END_COM_MAP()

    DECLARE_PROTECT_FINAL_CONSTRUCT()

    static HRESULT s_CreateInstance(__in IShellFolder *psf, PCUITEMID_CHILD pidl, REFIID riid, __deref_out void **ppv);

    HRESULT FinalConstruct()
    {
        return S_OK;
    }

    void FinalRelease()
    {
    }

    // IPropertyStoreFactory
    IFACEMETHODIMP GetPropertyStore(GETPROPERTYSTOREFLAGS flags, __in_opt IUnknown *punkFactory, REFIID riid, __deref_out void **ppv);
    IFACEMETHODIMP GetPropertyStoreForKeys(const PROPERTYKEY *rgKeys, UINT cKeys, GETPROPERTYSTOREFLAGS flags, REFIID riid, __deref_out void **ppv);

private:
    HRESULT _Initialize(__in IShellFolder *psf, PCUITEMID_CHILD pidl);
    HRESULT _ValidateFlags(GETPROPERTYSTOREFLAGS flags);
    HRESULT _CreatePropertyKeyArray(GETPROPERTYSTOREFLAGS flags, __inout CAtlArray<PROPERTYKEY> &rgKeys);

private:
    CComPtr<IShellFolder> _spsf;
    CComMemPtr<ITEMID_CHILD> _spidl;
};

HRESULT CRegPropertyStoreFactory::_Initialize(__in IShellFolder *psf, PCUITEMID_CHILD pidl)
{
    CComMemPtr<ITEMID_CHILD> spidl;
    HRESULT hr = CloneChildIDList(pidl, &spidl);
    if (SUCCEEDED(hr))
    {
        _spidl.Attach(spidl.Detach());
        _spsf = psf;
    }
    return hr;
}

HRESULT CRegPropertyStoreFactory::s_CreateInstance(__in IShellFolder *psf, PCUITEMID_CHILD pidl, REFIID riid, __deref_out void **ppv)
{
    *ppv = NULL;
    CComObjectWithRef<CRegPropertyStoreFactory> *psprpsf;
    HRESULT hr = CComObjectWithRef<CRegPropertyStoreFactory>::CreateInstance(&psprpsf);
    if (SUCCEEDED(hr))
    {
        hr = psprpsf->_Initialize(psf, pidl);
        if (SUCCEEDED(hr))
        {
            hr = psprpsf->QueryInterface(riid, ppv);
        }
        psprpsf->Release();
    }
    return hr;
}

STDAPI CreateRegistryPropertyStoreFactory(__in IShellFolder *psf, PCUITEMID_CHILD pidlChild, REFIID riid, __deref_out void **ppv)
{
    return CRegPropertyStoreFactory::s_CreateInstance(psf, pidlChild, riid, ppv);
}

// IPropertyStoreFactory
HRESULT CRegPropertyStoreFactory::_CreatePropertyKeyArray(GETPROPERTYSTOREFLAGS flags, __inout CAtlArray<PROPERTYKEY> &rgKeys)
{
    CComPtr<IShellFolder2> spsf;
    HRESULT hr = _spsf->QueryInterface(IID_PPV_ARGS(&spsf));
    if (SUCCEEDED(hr))
    {
        try
        {
            PROPERTYKEY key;
            for (UINT i = 0; SUCCEEDED(spsf->MapColumnToSCID(i, &key)) && SUCCEEDED(hr); i++)
            {
                SHCOLSTATEF scsf;
                hr = spsf->GetDefaultColumnState(i, &scsf);
                if (SUCCEEDED(hr))
                {
                    if (!(flags & GPS_FASTPROPERTIESONLY) || !(scsf & SHCOLSTATE_SLOW))
                    {
                        rgKeys.Add(key);
                    }
                }
            }
        }
        catch (CAtlException &e)
        {
            hr = e.m_hr;
        }
    }
    return hr;
}

STDMETHODIMP CRegPropertyStoreFactory::GetPropertyStore(GETPROPERTYSTOREFLAGS flags, __in_opt IUnknown *punkFactory, REFIID riid, __deref_out void **ppv)
{
    CAtlArray<PROPERTYKEY> rgKeys;
    HRESULT hr = _CreatePropertyKeyArray(flags, rgKeys);
    if (SUCCEEDED(hr))
    {
        hr = GetPropertyStoreForKeys(rgKeys.GetData(), (UINT)rgKeys.GetCount(), flags, riid, ppv);
    }
    return hr;
}

HRESULT CRegPropertyStoreFactory::_ValidateFlags(GETPROPERTYSTOREFLAGS flags)
{
    HRESULT hr = S_OK;
    if (((flags & GPS_HANDLERPROPERTIESONLY) && (flags & GPS_FASTPROPERTIESONLY)) ||
        ((flags & GPS_READWRITE) && (flags & GPS_BESTEFFORT)) ||       
        (((flags & (GPS_READWRITE | GPS_HANDLERPROPERTIESONLY)) && (flags & GPS_DELAYCREATION))))
    {
        hr = E_INVALIDARG;
    }
    return hr;
}

STDMETHODIMP CRegPropertyStoreFactory::GetPropertyStoreForKeys(const PROPERTYKEY *rgKeys, UINT cKeys, GETPROPERTYSTOREFLAGS flags, REFIID riid, __deref_out void **ppv)
{
    HRESULT hr = _ValidateFlags(flags);
    if (SUCCEEDED(hr))
    {
        CComPtr<IPropertyStore> sppropstore;
        hr = CRegPropertyStore::s_CreateInstance(rgKeys, cKeys, flags, _spsf, _spidl, IID_PPV_ARGS(&sppropstore));
        if (SUCCEEDED(hr))
        {
            hr = sppropstore->QueryInterface(riid, ppv);
        }
    }
    return hr;
}

